{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "accumarray like function\n",
      "======================================================================\n",
      "\n",
      "accum, a function like MATLAB's accumarray\n",
      "------------------------------------------\n",
      "\n",
      "NumPy doesn't include a function that is equivalent to MATLAB's\n",
      "`accumarray` function. The following function, `accum`, is close.\n",
      "\n",
      "Note that `accum` can handle n-dimensional arrays, and allows the data\n",
      "type of the result to be specified."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from itertools import product\n",
      "import numpy as np\n",
      "\n",
      "\n",
      "def accum(accmap, a, func=None, size=None, fill_value=0, dtype=None):\n",
      "    \"\"\"\n",
      "    An accumulation function similar to Matlab's `accumarray` function.\n",
      "\n",
      "    Parameters\n",
      "    ----------\n",
      "    accmap : ndarray\n",
      "        This is the \"accumulation map\".  It maps input (i.e. indices into\n",
      "        `a`) to their destination in the output array.  The first `a.ndim`\n",
      "        dimensions of `accmap` must be the same as `a.shape`.  That is,\n",
      "        `accmap.shape[:a.ndim]` must equal `a.shape`.  For example, if `a`\n",
      "        has shape (15,4), then `accmap.shape[:2]` must equal (15,4).  In this\n",
      "        case `accmap[i,j]` gives the index into the output array where\n",
      "        element (i,j) of `a` is to be accumulated.  If the output is, say,\n",
      "        a 2D, then `accmap` must have shape (15,4,2).  The value in the\n",
      "        last dimension give indices into the output array. If the output is\n",
      "        1D, then the shape of `accmap` can be either (15,4) or (15,4,1) \n",
      "    a : ndarray\n",
      "        The input data to be accumulated.\n",
      "    func : callable or None\n",
      "        The accumulation function.  The function will be passed a list\n",
      "        of values from `a` to be accumulated.\n",
      "        If None, numpy.sum is assumed.\n",
      "    size : ndarray or None\n",
      "        The size of the output array.  If None, the size will be determined\n",
      "        from `accmap`.\n",
      "    fill_value : scalar\n",
      "        The default value for elements of the output array. \n",
      "    dtype : numpy data type, or None\n",
      "        The data type of the output array.  If None, the data type of\n",
      "        `a` is used.\n",
      "\n",
      "    Returns\n",
      "    -------\n",
      "    out : ndarray\n",
      "        The accumulated results.\n",
      "\n",
      "        The shape of `out` is `size` if `size` is given.  Otherwise the\n",
      "        shape is determined by the (lexicographically) largest indices of\n",
      "        the output found in `accmap`.\n",
      "\n",
      "\n",
      "    Examples\n",
      "    --------\n",
      "    >>> from numpy import array, prod\n",
      "    >>> a = array([[1,2,3],[4,-1,6],[-1,8,9]])\n",
      "    >>> a\n",
      "    array([[ 1,  2,  3],\n",
      "           [ 4, -1,  6],\n",
      "           [-1,  8,  9]])\n",
      "    >>> # Sum the diagonals.\n",
      "    >>> accmap = array([[0,1,2],[2,0,1],[1,2,0]])\n",
      "    >>> s = accum(accmap, a)\n",
      "    array([9, 7, 15])\n",
      "    >>> # A 2D output, from sub-arrays with shapes and positions like this:\n",
      "    >>> # [ (2,2) (2,1)]\n",
      "    >>> # [ (1,2) (1,1)]\n",
      "    >>> accmap = array([\n",
      "            [[0,0],[0,0],[0,1]],\n",
      "            [[0,0],[0,0],[0,1]],\n",
      "            [[1,0],[1,0],[1,1]],\n",
      "        ])\n",
      "    >>> # Accumulate using a product.\n",
      "    >>> accum(accmap, a, func=prod, dtype=float)\n",
      "    array([[ -8.,  18.],\n",
      "           [ -8.,   9.]])\n",
      "    >>> # Same accmap, but create an array of lists of values.\n",
      "    >>> accum(accmap, a, func=lambda x: x, dtype='O')\n",
      "    array([[[1, 2, 4, -1], [3, 6]],\n",
      "           [[-1, 8], [9]]], dtype=object)\n",
      "    \"\"\"\n",
      "\n",
      "    # Check for bad arguments and handle the defaults.\n",
      "    if accmap.shape[:a.ndim] != a.shape:\n",
      "        raise ValueError(\"The initial dimensions of accmap must be the same as a.shape\")\n",
      "    if func is None:\n",
      "        func = np.sum\n",
      "    if dtype is None:\n",
      "        dtype = a.dtype\n",
      "    if accmap.shape == a.shape:\n",
      "        accmap = np.expand_dims(accmap, -1)\n",
      "    adims = tuple(range(a.ndim))\n",
      "    if size is None:\n",
      "        size = 1 + np.squeeze(np.apply_over_axes(np.max, accmap, axes=adims))\n",
      "    size = np.atleast_1d(size)\n",
      "\n",
      "    # Create an array of python lists of values.\n",
      "    vals = np.empty(size, dtype='O')\n",
      "    for s in product(*[range(k) for k in size]):\n",
      "        vals[s] = []\n",
      "    for s in product(*[range(k) for k in a.shape]):\n",
      "        indx = tuple(accmap[s])\n",
      "        val = a[s]\n",
      "        vals[indx].append(val)\n",
      "\n",
      "    # Create the output array.\n",
      "    out = np.empty(size, dtype=dtype)\n",
      "    for s in product(*[range(k) for k in size]):\n",
      "        if vals[s] == []:\n",
      "            out[s] = fill_value\n",
      "        else:\n",
      "            out[s] = func(vals[s])\n",
      "\n",
      "    return out"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Examples\n",
      "--------\n",
      "\n",
      "A basic example--sum the diagonals (with wrapping) of a 3 by 3 array:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from numpy import array, prod\n",
      "\n",
      "a = array([[1,2,3],[4,-1,6],[-1,8,9]])\n",
      "accmap = array([[0,1,2],[2,0,1],[1,2,0]])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "s = accum(accmap, a)\n",
      "s"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 3,
       "text": [
        "array([ 9,  7, 15])"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Accumulate using multiplication, going from a 3 by 3 array to 2 by 2\n",
      "array:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "accmap = array([\n",
      "             [[0,0],[0,0],[0,1]],\n",
      "             [[0,0],[0,0],[0,1]],\n",
      "             [[1,0],[1,0],[1,1]],\n",
      "         ])\n",
      "\n",
      "accum(accmap, a, func=prod, dtype=float)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 4,
       "text": [
        "array([[ -8.,  18.],\n",
        "       [ -8.,   9.]])"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Create an array of lists containing the values to be accumulated in each\n",
      "position in the output array:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "accum(accmap, a, func=lambda x: x, dtype='O')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 5,
       "text": [
        "array([[[1, 2, 4, -1], [3, 6]],\n",
        "       [[-1, 8], [9]]], dtype=object)"
       ]
      }
     ],
     "prompt_number": 5
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Use \\`accum\\` to arrange some values from a 1D array in a 2D array (note\n",
      "that using \\`accum\\` for this is overkill; fancy indexing would\n",
      "suffice):"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "subs = np.array([[k,5-k] for k in range(6)])\n",
      "subs"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 6,
       "text": [
        "array([[0, 5],\n",
        "       [1, 4],\n",
        "       [2, 3],\n",
        "       [3, 2],\n",
        "       [4, 1],\n",
        "       [5, 0]])"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "vals = array(range(10,16))\n",
      "accum(subs, vals)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 7,
       "text": [
        "array([[ 0,  0,  0,  0,  0, 10],\n",
        "       [ 0,  0,  0,  0, 11,  0],\n",
        "       [ 0,  0,  0, 12,  0,  0],\n",
        "       [ 0,  0, 13,  0,  0,  0],\n",
        "       [ 0, 14,  0,  0,  0,  0],\n",
        "       [15,  0,  0,  0,  0,  0]])"
       ]
      }
     ],
     "prompt_number": 7
    }
   ],
   "metadata": {}
  }
 ]
}